
#include <iostream>
#include <iomanip>
#include <queue>
#include "solution.h"

using namespace std;


solution::solution(solution* parent)
{
    evaluated = false;
    root = copySubtree(parent->getRoot());   
}

void solution::evaluate(dataset* trainSet, dataset* testSet)
{
    evaluated = true;
    for (int i = 0; i < 5; i++)
        for (int j = 0; j < 5; j++)
        {
            confusionTrain[i][j] = 0;
            confusionTest[i][j] = 0;
        }

    // Process all records in the training set
    assert(root != NULL);
    for (int i = 0; i < trainSet->getRecordCount(); i++)
    {
        root->processTrainingRecord(trainSet->getNthRecord(i));
    }

    // Delete any nodes (except root) that do not have the adequate size sets
    queue<node *> nodeQ;
    nodeQ.push(root);
    while(!nodeQ.empty())
    {
        node* tmp = nodeQ.front(); 
        nodeQ.pop();
  
        if (tmp->getLEObservationCount() < MIN_SPLIT_SIZE)
        {
            if (tmp->getLeft() != NULL)
            {
                delete tmp->getLeft();
                tmp->setLeft(NULL);
            }
        }
        if (tmp->getGTObservationCount() < MIN_SPLIT_SIZE)
        {
            if (tmp->getRight() != NULL)
            {
                delete tmp->getRight();
                tmp->setRight(NULL);
            }
        }
        
        if (tmp->getLeft() != NULL) 
            nodeQ.push(tmp->getLeft());

        if (tmp->getRight() != NULL) 
            nodeQ.push(tmp->getRight());
    }
        

    // Collect the stats on the training set by visiting the leaves
    nodeQ.push(root);
    while(!nodeQ.empty())
    {
        node* tmp = nodeQ.front(); 
        nodeQ.pop();

        if (tmp->getLeft() != NULL) 
            nodeQ.push(tmp->getLeft());
        else
        {
            int pred = tmp->getLEWinner();
            for (int i = 0; i < 5; i++)
            {
                 confusionTrain[i][pred] += tmp->getLEObservationCount(i);
            }
        }
        if (tmp->getRight() != NULL)
            nodeQ.push(tmp->getRight());
        else
        {
            int pred = tmp->getGTWinner();
            for (int i = 0; i < 5; i++)
            {
                 confusionTrain[i][pred] += tmp->getGTObservationCount(i);
            }
        }
    }

    for (int i = 0; i < 5; i++)
    {
        TrainPD[i] = (double) confusionTrain[i][i] / 
            (confusionTrain[i][0] + confusionTrain[i][1] + confusionTrain[i][2] + 
             confusionTrain[i][3] + confusionTrain[i][4]);
        double alarms = confusionTrain[0][i] + confusionTrain[1][i] +
             confusionTrain[2][i] + confusionTrain[3][i] + confusionTrain[4][i];
        TrainFAR[i] = (alarms - confusionTrain[i][i]) / alarms;
    }

    for (int i = 0; i < testSet->getRecordCount(); i++)
    {
        record* tmpRecord = testSet->getNthRecord(i);
        int pred = root->evaluateTestRecord(tmpRecord);

        confusionTest[tmpRecord->getResult()][pred]++;
    }

    for (int i = 0; i < 5; i++)
    {
        TestPD[i] = (double) confusionTest[i][i] /
            (confusionTest[i][0] + confusionTest[i][1] + confusionTest[i][2] +
             confusionTest[i][3] + confusionTest[i][4]);
        double alarms = confusionTest[0][i] + confusionTest[1][i] +
             confusionTest[2][i] + confusionTest[3][i] + confusionTest[4][i];
        TestFAR[i] = (alarms - confusionTest[i][i]) / alarms;
    }

}

void solution::printTree(node* tmp, int depth)
{
    cout << "T ";
    for (int n = 0; n < depth; n++)
        cout << "  ";

    cout << "if " << dataset::getFieldName(tmp->getField());
    if (dataset::isFieldCategorical(tmp->getField()))
    {
        cout << " in {";
        for (set<int>::iterator it = tmp->getValues().begin(); it != tmp->getValues().end(); ++it)
        {
            cout << " '" << dataset::getFieldValue(tmp->getField(), *it) << "'";
        }
        cout << "} \n";
    }
    else
    {
        cout << " <= " << tmp->getValue() << endl;
    }

    if (tmp->getLeft() == NULL)
    {
        cout << "T ";
        for (int n = 0; n <= depth; n++)
            cout << "  ";
        cout << dataset::getAttackName(tmp->getLEWinner()) << endl;
    }
    else
        printTree(tmp->getLeft(), depth + 1);

    cout << "T ";
    for (int n = 0; n < depth; n++)
        cout << "  ";
    cout << "else\n";

    if (tmp->getRight() == NULL)
    {
        cout << "T ";
        for (int n = 0; n <= depth; n++)
            cout << "  ";
        cout << dataset::getAttackName(tmp->getGTWinner()) << endl;
    }
    else
        printTree(tmp->getRight(), depth + 1);
}

void solution::display()
{
    cout << "+ Set\tFitness\t1PD\t1FAR\t2PD\t2FAR\t3PD\t3FAR\t4PD\t4FAR\n";
    cout << "+ Train\t" << getTrainFitness();
    for (int i = 1; i < 5; i++)
         cout << "\t" <<  TrainPD[i] << "\t" << TrainFAR[i]; 
    cout << endl;

    cout << "+ Test\t" << getTestFitness();
    for (int i = 1; i < 5; i++)
         cout << "\t" << TestPD[i] << "\t" << TestFAR[i]; 
    cout << endl;

    cout << "+ Training Set Confusion Matrix\n";
    cout << "+ " << setw(12) << "Normal"
                 << setw(12) << "Probe"
                 << setw(12) << "DoS"
                 << setw(12) << "User2Root"
                 << setw(12) << "Remote2Local\n";
    for (int i = 0; i < 5; i++)
    {
        cout << "+ ";
        for (int j = 0; j < 5; j++)
            cout << setw(12) << confusionTrain[i][j];
        cout << endl;
    }
        
    cout << "+ Test Set Confusion Matrix\n";
    cout << "+ " << setw(12) << "Normal"
                 << setw(12) << "Probe"
                 << setw(12) << "DoS"
                 << setw(12) << "User2Root"
                 << setw(12) << "Remote2Local\n";
    for (int i = 0; i < 5; i++)
    {
        cout << "+ ";
        for (int j = 0; j < 5; j++)
            cout << setw(12) << confusionTest[i][j];
        cout << endl;
    }


    printTree(root, 0);
}
